#pragma once

#include <iostream>

using std::cout;
using std::endl;

template <typename T, typename Compare = std::less<int>>
class RB_Tree
{
    friend void test1();
    friend void test2();

private:
    enum class Color
    {
        RED,
        BLACK
    };

    struct Node
    {
        Node(T val = T(), Node *left = nullptr, Node *right = nullptr, Node *parent = nullptr, Color color = Color::BLACK)
            : val_(val), left_(left), right_(right), parent_(parent), color_(color) {}
        T val_;
        Node *left_;
        Node *right_;
        Node *parent_;
        Color color_;
    };

    Node *root_;
    Compare comp_;

    void rightRotate(Node *node)
    {
        Node *child = left(node);
        // 如果当前节点右父节点
        if (parent(node))
        {
            // 如果当前节点是父节点的左孩子
            if (parent(node)->left_ == node)
            {
                parent(node)->left_ = child;
            }
            else
            {
                parent(node)->right_ = child;
            }
        }
        child->parent_ = parent(node);

        node->left_ = child->right_;
        if (right(child))
        {
            child->right_->parent_ = node;
        }

        node->parent_ = child;
        child->right_ = node;

        if (!parent(child))
        {
            root_ = child;
        }
    }

    void leftRotate(Node *node)
    {
        Node *child = right(node);
        // 如果当前节点右父节点
        if (parent(node))
        {
            // 如果当前节点是父节点的左孩子
            if (parent(node)->left_ == node)
            {
                parent(node)->left_ = child;
            }
            else
            {
                parent(node)->right_ = child;
            }
        }
        child->parent_ = parent(node);

        node->right_ = child->left_;
        if (left(child))
        {
            child->left_->parent_ = node;
        }

        node->parent_ = child;
        child->left_ = node;

        if (!parent(child))
        {
            root_ = child;
        }
    }

    Color color(Node *node) const
    {
        if (!node)
            return Color::BLACK;
        return node->color_;
    }

    Node *parent(Node *node) const
    {
        return node->parent_;
    }

    Node *left(Node *node) const
    {
        return node->left_;
    }

    Node *right(Node *node) const
    {
        return node->right_;
    }

    void setColor(Node *node, Color color)
    {
        node->color_ = color;
    }

    void fixAfterInsert(Node *node)
    {
        // 一直调整到父节点为黑色
        while (color(parent(node)) == Color::RED)
        {
            // 如果当前插入的节点在爷爷节点的左子树中
            if (left(parent(parent(node))) == parent(node))
            {
                Node *uncle = right(parent(parent(node)));
                // 情况一
                if (Color::RED == color(uncle))
                {
                    setColor(parent(node), Color::BLACK);
                    setColor(parent(parent(node)), Color::RED);
                    setColor(uncle, Color::BLACK);
                }
                else
                {
                    // 情况三：插入节点在父节点的右侧
                    if (right(parent(node)) == node)
                    {
                        leftRotate(parent(node));
                        // 更新node节点
                        node = left(node);
                    }
                    // 情况二：插入节点在父节点的左侧

                    setColor(parent(node), Color::BLACK);
                    setColor(parent(parent(node)), Color::RED);
                    rightRotate(parent(parent(node)));
                    // 表示调整已经完成
                    break;
                }
            }
            else
            {
                Node *uncle = left(parent(parent(node)));
                // 情况一
                if (Color::RED == color(uncle))
                {
                    setColor(parent(node), Color::BLACK);
                    setColor(parent(parent(node)), Color::RED);
                    setColor(uncle, Color::BLACK);
                }
                else
                {
                    // 情况三：插入节点在父节点的右侧
                    if (left(parent(node)) == node)
                    {
                        rightRotate(parent(node));
                        // 更新node节点
                        node = right(node);
                    }
                    // 情况二：插入节点在父节点的左侧

                    setColor(parent(node), Color::BLACK);
                    setColor(parent(parent(node)), Color::RED);
                    leftRotate(parent(parent(node)));
                    // 表示调整已经完成
                    break;
                }
            }

            node = parent(parent(node));
        }

        // 如果当前节点没有父节点，说明是root，
        setColor(root_, Color::BLACK);
    }

    void insert(Node *node, const T &val)
    {
        Node *parent = nullptr;
        while (node)
        {
            parent = node;
            if (!comp_(node->val_, val))
            {
                node = node->left_;
            }
            else if (comp_(node->val_, val))
            {
                node = node->right_;
            }
            else
            {
                return;
            }
        }

        Node *cur = new Node(val, nullptr, nullptr, parent, Color::RED);
        if (parent->val_ > val)
        {
            parent->left_ = cur;
        }
        else
        {
            parent->right_ = cur;
        }

        if (color(parent) == Color::RED)
        {
            fixAfterInsert(cur);
        }
    }

    void fixAfterRemove(Node *node)
    {
        while (node != root_ && color(node) == Color::BLACK)
        {
            // 左子树情况
            if (left(parent(node)) == node)
            {
                Node *bro = right(parent(node));
                // 如果是情况四
                if (color(bro) == Color::RED)
                {
                    setColor(bro, Color::BLACK);
                    setColor(parent(bro), Color::RED);
                    leftRotate(parent(bro));
                    bro = right(parent(node));
                }

                // 如果是情况三
                if (color(left(bro)) == Color::BLACK && color(right(bro)) == Color::BLACK)
                {
                    setColor(bro, Color::RED);
                    node = parent(node);
                }
                else
                {
                    // 情况二
                    if (color(right(bro)) != Color::RED)
                    {
                        setColor(bro, Color::RED);
                        setColor(left(bro), Color::BLACK);
                        rightRotate(bro);
                        bro = right(parent(node));
                    }

                    if (color(right(bro)) == Color::RED)
                    {
                        setColor(bro, Color::BLACK);
                        setColor(parent(bro), Color::BLACK);
                        setColor(right(bro), Color::BLACK);
                        leftRotate(parent(bro));
                        break;
                    }
                }
            }
            else
            {
                Node *bro = left(parent(node));
                // 如果是情况四
                if (color(bro) == Color::RED)
                {
                    setColor(bro, Color::BLACK);
                    setColor(parent(bro), Color::RED);
                    rightRotate(parent(bro));
                    bro = left(parent(node));
                }

                // 如果是情况三
                if (color(right(bro)) == Color::BLACK && color(left(bro)) == Color::BLACK)
                {
                    setColor(bro, Color::RED);
                    node = parent(node);
                }
                else
                {
                    // 情况二
                    if (color(left(bro)) != Color::RED)
                    {
                        setColor(bro, Color::RED);
                        setColor(right(bro), Color::BLACK);
                        leftRotate(bro);
                        bro = left(parent(node));
                    }

                    if (color(left(bro)) == Color::RED)
                    {
                        setColor(bro, Color::BLACK);
                        setColor(parent(bro), Color::RED);
                        setColor(left(bro), Color::BLACK);
                        rightRotate(parent(bro));
                        break;
                    }
                }
            }
        }
        setColor(node, Color::BLACK);
    }

    void remove(Node *node, const T &val)
    {
        while (node)
        {
            if (node->val_ > val)
            {
                node = node->left_;
            }
            else if (node->val_ < val)
            {
                node = node->right_;
            }
            else
            {
                break;
            }
        }

        // 没有要删除的元素
        if (!node)
        {
            return;
        }
        else
        {
            // 如果左右孩子都不为空
            if (node->left_ && node->right_)
            {
                Node *pre = node->left_;
                while (pre->right_)
                {
                    pre = pre->right_;
                }
                node->val_ = pre->val_;
                node = pre;
            }

            Node *child = node->left_;
            if (!child)
            {
                child = node->right_;
            }

            // 如果要删除的节点是最后一个root节点
            if (!child && !parent(node))
            {
                delete node;
                root_ = nullptr;
                return;
            }
            else
            { // 如果当前节点有孩子节点，并且要删除的节点有父节点
                if (child)
                {
                    child->parent_ = parent(node);
                    if (left(parent(node)) == node)
                    {
                        node->parent_->left_ = child;
                    }
                    else
                    {
                        node->parent_->right_ = child;
                    }

                    if (color(node) == Color::BLACK)
                    {
                        fixAfterRemove(child);
                    }
                }
                else
                {
                    if (color(node) == Color::BLACK)
                    {
                        fixAfterRemove(node);
                    }
                    if (left(parent(node)) == node)
                    {
                        node->parent_->left_ = nullptr;
                    }
                    else
                    {
                        node->parent_->right_ = nullptr;
                    }
                }
            }
        }

        delete node;
    }

public:
    RB_Tree() : root_(nullptr) {}
    ~RB_Tree() {}

    void remove(const T &val)
    {
        remove(root_, val);
    }

    void insert(const T &val)
    {
        if (!root_)
        {
            root_ = new Node(val);
            return;
        }
        insert(root_, val);
    }
};